{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 27/27 [00:00<00:00, 12543.89it/s]\n"
     ]
    }
   ],
   "source": [
    "import  os \n",
    "from os.path import join, isdir, isfile, exists\n",
    "from os import makedirs\n",
    "\n",
    "import cv2\n",
    "from glob import glob\n",
    "\n",
    "import json\n",
    "\n",
    "\n",
    "import numpy as np \n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "import shutil\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class SeqDataloader:\n",
    "\n",
    "\n",
    "    def __init__(self, path):\n",
    "        self.path  = path\n",
    "        \n",
    "        self.img_paths = sorted(glob(join(path,'img1','*.jpg')))\n",
    "\n",
    "        self.gt = np.loadtxt(join(path,'gt','gt.txt'),delimiter=',',dtype=np.int32)\n",
    "\n",
    "        \n",
    "\n",
    "    def __len__(self,):\n",
    "        return len(self.img_paths)\n",
    "\n",
    "    def split(self,ratio = 0.85):\n",
    "\n",
    "        \n",
    "        target_set_path = self.path.replace('train','test')\n",
    "        \n",
    "        def make_dir(path):\n",
    "            if not exists(path):\n",
    "                os.makedirs(path)\n",
    "        make_dir(target_set_path)\n",
    "\n",
    "        target_gt_path = join(target_set_path,'gt')\n",
    "        target_img_path = join(target_set_path,'img1')\n",
    "        make_dir(target_img_path)\n",
    "        make_dir(target_gt_path)\n",
    "        target_gt_path = join(target_gt_path, 'gt.txt')\n",
    "        \n",
    "        \n",
    "        \n",
    "        train_border = int(self.__len__() * ratio)\n",
    "\n",
    "        \n",
    "        for img_path in tqdm(self.img_paths[train_border:]):\n",
    "            shutil.move(img_path,join(target_img_path,img_path.split('/')[-1]))\n",
    "\n",
    "\n",
    "\n",
    "        img_idx = self.gt[:,0].astype(np.int32)\n",
    "        testset_gt = self.gt[img_idx>train_border,:]\n",
    "        trainset_gt = self.gt[img_idx<=train_border,:]\n",
    "\n",
    "\n",
    "        def transform(gt_labels):\n",
    "            new_one = []\n",
    "            for line in gt_labels.tolist():            \n",
    "                line[0] = str(line[0]).rjust(6, '0')\n",
    "                new_one.append(line)\n",
    "            return new_one\n",
    "\n",
    "        testset_gt = transform(testset_gt.astype(np.str0))\n",
    "        trainset_gt = transform(trainset_gt.astype(np.str0))\n",
    "\n",
    "        def write_file(data,path):\n",
    "            with open(path,'w') as f :\n",
    "                f.write('\\n'.join([','.join(x) for x in data]))\n",
    "        \n",
    "        # np.savetxt(target_gt_path, testset_gt, delimiter=',',fmt='%d')\n",
    "        # np.savetxt(join(self.path,'gt','gt.txt'), trainset_gt, delimiter=',',fmt='%d')\n",
    "\n",
    "        write_file(testset_gt,target_gt_path)\n",
    "        write_file(trainset_gt,join(self.path,'gt','gt.txt'))\n",
    "\n",
    "\n",
    "    \n",
    "    '''\n",
    "    description: 797 for ant09\n",
    "    param {*} self\n",
    "    param {*} step\n",
    "    return {*}\n",
    "    '''\n",
    "    def reorganize_testset(self,step = 797):\n",
    "        target_set_path = self.path.replace('train','test')\n",
    "\n",
    "        target_gt_path = join(target_set_path,'gt','gt.txt')\n",
    "        target_img_path = join(target_set_path,'img1')\n",
    "\n",
    "        gt = np.loadtxt(target_gt_path,delimiter = ',', dtype = np.str0).astype(np.int32)\n",
    "\n",
    "        gt[:,0] -= step\n",
    "\n",
    "        gt = gt.astype(np.str0)\n",
    "        new_one = []\n",
    "        for line in gt.tolist():            \n",
    "            line[0] = str(line[0]).rjust(6, '0')\n",
    "            new_one.append(line)\n",
    "        \n",
    "        with open(target_gt_path,'w') as f :\n",
    "            f.write('\\n'.join([','.join(x) for x in new_one]))\n",
    "\n",
    "\n",
    "        \n",
    "        img_paths = sorted(glob(target_img_path+'/*.jpg'))\n",
    "\n",
    "        # print(new_one)\n",
    "        # print(img_paths)\n",
    "        for idx, img in enumerate(img_paths):\n",
    "            # img_name = img.split('/')[-1].split('.')[0]\n",
    "            # print(img_name)\n",
    "\n",
    "            new_name = str(idx+1).rjust(6, '0') + '.jpg'\n",
    "            \n",
    "            new_name = '/'.join(img.split('/')[:-1] + [new_name])\n",
    "            # print(new_name,img)\n",
    "            shutil.move(img,new_name)\n",
    "\n",
    "            \n",
    "            \n",
    "    \n",
    "    \n",
    "        \n",
    "\n",
    "# src_seq = '/data/xusc/exp/topictrack-bee/data/antmove/train/ant09'\n",
    "src_seq = '/data/xusc/exp/topictrack-bee/data/beedance/train/bee03'\n",
    "seq_dataloader  = SeqDataloader(src_seq)\n",
    "\n",
    "seq_dataloader.split()\n",
    "seq_dataloader.reorganize_testset(step = 147)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([[  1,   1, 211, ...,   1,   1,   1],\n",
       "        [  1,   2, 249, ...,   1,   1,   1],\n",
       "        [  1,   3, 402, ...,   1,   1,   1],\n",
       "        ...,\n",
       "        [938,  46, 156, ...,   1,   1,   1],\n",
       "        [938,  13, 237, ...,   1,   1,   1],\n",
       "        [938,  11, 209, ...,   1,   1,   1]], dtype=int32),\n",
       " 938)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "seq_dataloader.img_paths\n",
    "seq_dataloader.gt, seq_dataloader.__len__()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "yang_real",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
